# -*- coding: utf-8 -*-
"""
Created on Mon Oct 17 18:51:06 2022

@author: ZJH
"""


# --------- Core Structure Base Class ---------


class BaseCore:
    """

    Base class for RTMQ core structure,
    defines constants and registers common to the RTMQ framework.

    Constants
    ----------
    OPC : Type-A instruction opcode list. \n
    MGW : Magic word of RTLink frame. \n
    RSV : Reserved registers (core registers & system peripherals).

    Attributes
    ----------
    REG : List of register names.

    """
    OPC = ["ADD", "AND", "XOR",
           "CLU", "CLS", "CEQ",
           "SGN", "SNE", "SMK", "MOV",
           "SLL", "SLA", "SLC", "REV"]
    MGW = 0xA
    RSV = ["NUL", "PTR", "WCK", "TIM", "TRG", "ERR",
           "ICA", "ICD", "DCA", "DCD", "RTS", "DTS",
           "RTCF", "RTBC", "RTD0"]

    def __init__(self, regs):
        """

        Parameters
        ----------
        regs : list of strings
            List of all other registers.

        Returns
        -------
        None.

        """
        self.REG = self.RSV + regs
        self.alu_opc = dict()
        for i in range(len(self.OPC)):
            self.alu_opc[self.OPC[i]] = i
        self.regs = dict()
        for i in range(len(self.REG)):
            self.regs[self.REG[i]] = i


# Global instance used to implement basic operations
C_BASE = BaseCore([])


# --------- Assemble Functions ---------


def bit_concat(*args):
    """

    Bit fields concatenation.

    Parameters
    ----------
    *args : 2-tuples as (V, W)
        Each argument corresponds to a bit field with value V and width W.
        Bit fields are concatenated from high to low.
        NOTE: V will be cropped if not less than (2 ** W)

    Returns
    -------
    dat : integer
        Concatenated value.

    Example
    -------
    bit_concat((0b11, 2), (0b10, 3), (0b11, 1)) == 0b110101

    """
    dat = 0
    ln = len(args)
    for i in range(ln - 1):
        dat = dat + args[i][0] % (2 ** args[i][1])
        dat = dat * (2 ** args[i + 1][1])
    dat = dat + args[ln - 1][0] % (2 ** args[ln - 1][1])
    return dat


def asm_instr(ins, core=None, label=dict()):
    """

    Convert an assembly instruction to machine code.

    Parameters
    ----------
    ins : string
        Assembly instruction.
    core : <CoreBase>, optional
        Destination core structure descriptor. The default is None.
    label : dict, optional
        Dict of jump labels. The default is dict().

    Returns
    -------
    code : integer
        Machine code.

    """

    def cnv_opd(idx, field):
        try:
            opd = c[idx]
            if opd[0] == "!":
                opd = opd[1:]
                inv = 1
            else:
                inv = 0
            if opd in core.REG:
                res = (core.regs[opd], 0, inv)
            elif opd in label.keys():
                res = (label[opd], 1, 0)
            else:
                res = (int(opd, 0) & 0xFFFF_FFFF, 1, 0)
            if (field == "RD" and res[1] == 1) or \
               (c[0] in {"LDL", "LDH"} and field == "R0" and res[1] == 0) or \
               (c[0] == "MOV" and field == "R0" and res[1] == 1):
                raise
        except:
            raise SyntaxError(err_msg + f"Invalid {field}.")
        return res

    ins = ins.strip().upper()
    core = C_BASE if core is None else core
    err_msg = f"Error in compiling '{ins}': "
    tmp = ins.split(" ")
    c = [s for s in tmp if s != ""]
    if len(c) < 3 or c[1] not in {"-", "H"} or c[2] not in {"-", "F"} :
        raise SyntaxError(err_msg + "Invalid HLD/FCT flag.")
    H = int(c[1] == "H")
    F = int(c[2] == "F")
    if c[0] == "NOP":
        return bit_concat((0, 2), (H, 1), (F, 1), (0, 28))
    RD, tmp, nRD = cnv_opd(3, "RD")
    R0, iR0, nR0 = cnv_opd(4, "R0")
    if c[0] == "LDH":
        return bit_concat((0, 2), (H, 1), (F, 1), (0, 4), (RD, 8),
                          (R0 >> 20, 16))
    elif c[0] == "LDL":
        return bit_concat((1, 2), (H, 1), (F, 1), (R0 >> 16, 4),
                          (RD, 8), (R0, 16))
    if c[0] not in core.OPC:
        raise SyntaxError(err_msg + "Invalid instruction.")
    opc = core.alu_opc[c[0]]
    if c[0] == "MOV":
        return bit_concat((2 + nRD, 2), (H, 1), (F, 1), (opc, 4),
                          (RD, 8), (4, 3), (nR0, 1), (0, 4), (R0, 8))
    else:
        if c[0] == "REV":
            R1, iR1, nR1 = (0, 0, 0)
        else:
            R1, iR1, nR1 = cnv_opd(5, "R1")
        if iR0 == 1 and (0x000000_7F < R0 < 0xFFFFFF_C0):
            raise ValueError(err_msg + "R0 immediate out of bound.")
        elif iR0 == 0 and R0 > 0x3F:
            raise ValueError(err_msg + "R0 address out of bound.")
        if iR1 == 1 and (0x000000_7F < R1 < 0xFFFFFF_C0):
            raise ValueError(err_msg + "R1 immediate out of bound.")
        elif iR1 == 0 and R1 > 0x3F:
            raise ValueError(err_msg + "R1 address out of bound.")
        t0 = (R0 >> 6) % 2 if iR0 else nR0
        t1 = (R1 >> 6) % 2 if iR1 else nR1
        return bit_concat((2 + nRD, 2), (H, 1), (F, 1), (opc, 4),
                          (RD, 8), (iR0, 1), (iR1, 1),
                          (t0, 1), (t1, 1), (R0, 6), (R1, 6))


def asm_program(prg, core=None):
    """

    Convert an assembly program to machine code.

    Parameters
    ----------
    prg : multi-line string OR list of strings
        Assembly program.
    core : <CoreBase>, optional
        Destination core structure descriptor. The default is None.

    Returns
    -------
    code : list of integers
        Machine code.

    """
    tmp = prg.splitlines() if isinstance(prg, str) else prg
    label = dict()
    lnum = 0
    lines = []
    for s in tmp:
        s = s.strip().upper()
        if s == "" or s[0] == "%":
            continue
        elif s[0] == "#" and s[-1] == ":":
            label[s[:-1]] = lnum
        else:
            lines += [s]
            lnum += 1
    code = [0] * lnum
    for i in range(lnum):
        code[i] = asm_instr(lines[i], core, label)
    return code


# --------- Basic Debug Functions ---------


def bit_split(dat, wids):
    """

    Split a number into bit fields.

    Parameters
    ----------
    dat : integer
        Number to be splitted.
    wids : list OR tuple of integers
        Width of each bit field.

    Returns
    -------
    bf : list of integers
        Values of bit fields.

    Example
    -------
    bit_split(0b110101, (2, 3, 1)) == [0b11, 0b10, 0b1]

    """
    bf = []
    for i in range(1, len(wids)):
        t = 2 ** wids[-i]
        bf += [dat % t]
        dat = dat // t
    bf += [dat % (2 ** wids[0])]
    bf.reverse()
    return bf

def raw_form(code):
    """

    Display the machine code program in raw form.

    Parameters
    ----------
    code : list of integers
        Program in machine code.

    Returns
    -------
    prg : multi-line string
        Program in raw form.

    """
    prg = ""
    for ins in code:
        ith, itl, ff, opc, rd, i0, i1, n0, n1, rx = \
            bit_split(ins, (1, 1, 2, 4, 8, 1, 1, 1, 1, 12))
        c = bit_concat((ith, 1), (itl, 1), (ff, 2))
        if ith == 0:
            imm = bit_concat((opc, 4), (i0, 1), (i1, 1),
                             (n0, 1), (n1, 1), (rx, 12))
            if itl == 0:
                i = f"{imm:03X}_00000"
            else:
                i = f"000_{imm:05X}"
            s = f"C: {c:04b}      RD: {rd:02X}         I: {i}"
        else:
            r0, r1 = (0, rx) if opc == 9 else bit_split(rx, (6, 6))
            if i0 == 1:
                r0 = bit_concat((n0, 1), (r0, 6))
                n0 = "_"
            if i1 == 1:
                r1 = bit_concat((n1, 1), (r1, 6))
                n1 = "_"
            rf = f"{i0}{i1}{n0}{n1}"
            s = f"C: {c:04b} O: {opc:01X} RD: {rd:02X} F: {rf} R0: {r0:02X} R1: {r1:02X}"
        prg += (s + "\n")
    return prg


def asm_form(code, core=None):
    """

    Convert machine code program back to assembly program.

    Parameters
    ----------
    code : list of integers
        Program in machine code.
    core : <CoreBase>, optional
        Destination core structure descriptor. The default is None.

    Returns
    -------
    prg : multi-line string
        Assembly program.

    """
    prg = ""
    core = C_BASE if core is None else core
    for ins in code:
        ith, itl, h, f, opc, rd, i0, i1, n0, n1, rx = \
            bit_split(ins, (1, 1, 1, 1, 4, 8, 1, 1, 1, 1, 12))
        cod = ("LDH", "LDL")[itl] if ith == 0 else core.OPC[opc]
        if cod == "LDH" and rd == 0:
            cod = "NOP"
        hld = ("-", "H")[h]
        fct = ("-", "F")[f]
        head = f"{cod} {hld} {fct}"
        rds = f"{core.REG[rd]:4}"
        if ith == 0:
            imm = bit_concat((opc, 4), (i0, 1), (i1, 1),
                             (n0, 1), (n1, 1), (rx, 12))
            if itl == 0:
                ims = f" 0x{imm:03X}_00000"
            else:
                ims = f" 0x000_{imm:05X}"
            if cod == "NOP":
                asm = head
            else:
                asm = f"{head}  {rds} {ims}"
        else:
            ird = (" ", "!")[itl]
            ir0 = (" ", "!")[n0]
            ir1 = (" ", "!")[n1]
            r0, r1 = (0, rx) if opc == 9 else bit_split(rx, (6, 6))
            r0i = f" 0x{bit_concat((n0, 1), (r0, 6)):02X}"
            r0s = f"{ir0}{core.REG[r0]}" if i0 == 0 else r0i
            r1i = f" 0x{bit_concat((n1, 1), (r1, 6)):02X}"
            r1s = f"{ir1}{core.REG[r1]}" if i1 == 0 else r1i
            if cod == "MOV":
                asm = f"{head} {ird}{rds} {r1s}"
            elif cod == "REV":
                asm = f"{head} {ird}{rds} {r0s}"
            else:
                asm = f"{head} {ird}{rds} {r0s} {r1s}"
        prg += (asm + "\n")
    return prg


def sim_form(code):
    """

    Generate Instruction Cache initialization data for behavioral simualtion.

    Parameters
    ----------
    code : list of integers
        Program in machine code.

    Returns
    -------
    prg : multi-line string
        Instruction Cache initialization data.

    """
    prg = ""
    for ins in code:
        prg += f"{ins:08X}\n"
    return prg


# --------- Basic RTLink Functions ---------


def _unpack_addr(addr):
    s = hex(addr).lstrip("0x")
    if "0" in s:
        raise RuntimeError(f"Invalid node address: 0x{addr:X}.")
    res = []
    for c in s:
        res += [int(c, 16)]
    return res


def _pack_frame(hop, payload, typ=1):
    frm = bit_concat((C_BASE.MGW, 4), (hop, 3), (typ, 1), (payload, 32))
    return frm.to_bytes(5, "big")


def _unpack_frame(frm):
    dat = int.from_bytes(frm, "big")
    mgw, hop, typ, pld = bit_split(dat, (4, 3, 1, 32))
    return pld, (mgw, hop, typ)


def get_payloads(buf):
    """

    Unpack RTLink frames from a device to get payloads.

    Parameters
    ----------
    buf : bytes-like iterable
        Byte stream of the frames.

    Returns
    -------
    payloads : list of integers
        Payloads of the frames.

    """
    ln = len(buf) // 5
    payloads = [0] * ln
    for i in range(ln):
        d, nil = _unpack_frame(buf[i*5:(i+1)*5])
        payloads[i] = d
    return payloads


def send_to(addr, payloads, typ=1):
    """

    Send data / instruction frames to a node.

    Parameters
    ----------
    addr : integer
        Destination node address.
    payloads : list of integers
        Payloads of frames, each element corresponds to a frame.
    typ : integer
        Type of frames.
        0: data frame; 1: instruction frame;
        The default is 1.

    Returns
    -------
    buf : bytearray
        Byte stream of the frames.

    """
    hop = len(_unpack_addr(addr))
    ln = len(payloads)
    buf = bytearray(ln * 5)
    for i in range(ln):
        buf[i*5:(i+1)*5] = _pack_frame(hop, payloads[i], typ)
    return buf


def node_exec(addr, prg, core=None):
    """

    Run a script on a node using configuration override.

    Parameters
    ----------
    addr : integer
        Destination node address.
    prg : multi-line string OR list of strings
        Assembly script to be run.
    core : <CoreBase>, optional
        Destination core structure descriptor. The default is None.

    Returns
    -------
    buf : bytearray
        Byte stream of the script.

    """
    code = asm_program(prg, core)
    return send_to(addr, code, 1)


def suspend_node(addr):
    """

    Generate a script to suspend the destination node.

    Parameters
    ----------
    addr : integer
        Destination node address.

    Returns
    -------
    buf : bytearray
        Byte stream of the script.

    """
    return node_exec(addr, "SMK - -  ERR  1  1")


def resume_node(addr):
    """

    Generate a script to resume the destination node.

    Parameters
    ----------
    addr : integer
        Destination node address.

    Returns
    -------
    buf : bytearray
        Byte stream of the script.

    """
    return node_exec(addr, "SMK - -  ERR  0  1")


def reset_node(addr):
    """

    Generate a script to reset the Timer, TrigMgr & ErrorMgr of a node,
    and then jump to address 0.

    Parameters
    ----------
    addr : integer
        Destination node address.

    Returns
    -------
    buf : bytearray
        Byte stream of the script.

    """
    asm = """
        LDL H F  PTR  0
        MOV H -  TRG  TRG
        LDL H -  TIM  0
        AND - - !ERR  ERR  NUL
        """
    return node_exec(addr, asm)


def download_to(addr, payloads, core=None, typ=1, start=0):
    """

    Generate a script to download data / program to the destination node.

    Parameters
    ----------
    addr : integer
        Destination node address.
    payloads : list of integers
        data / program in machine code.
    core : <CoreBase>, optional
        Destination core structure descriptor. The default is None.
    typ : integer, optional
        Type of payloads.
        0: data (to Data Cache); 1: program (to Instruction Cache);
        The default is 1.
    start : integer, optional
        Start address in the cache. The default is 0.

    Returns
    -------
    buf : bytearray
        Byte stream of the script.

    """
    ln = len(payloads)
    scr = [0] * (ln * 3)
    rga = ("DCA", "ICA")[typ]
    rgd = ("DCD", "ICD")[typ]
    for i in range(ln):
        scr[i*3+0] = f"LDL - -  {rga}  {i + start}"
        scr[i*3+1] = f"LDH - -  {rgd}  {payloads[i]}"
        scr[i*3+2] = f"LDL - -  {rgd}  {payloads[i]}"
    scr = [f"LDH - -  {rga}  0"] + scr
    return node_exec(addr, scr, core)


def reach_node(addr, ssp=True):
    """

    Construct a link between the computer and the destination node.
    So that the node can be controlled directly.

    Parameters
    ----------
    addr : integer
        Destination node address.
    ssp : bool, optional
        Whether to suspend the cores in the route. The default is True.

    Returns
    -------
    buf : bytearray
        Byte stream of the script.

    """
    nodes = _unpack_addr(addr)
    hop = len(nodes)
    buf = bytearray()
    for i in range(hop):
        t_adr = addr >> ((hop - i) * 4)
        if ssp:
            buf += suspend_node(t_adr)
        rtcf = bit_concat((nodes[i], 4), (0, 12), ((1 << nodes[i]) + 1, 16))
        scr = [f"LDH - -  RTCF  {rtcf}",
               f"LDL - -  RTCF  {rtcf}"]
        buf += node_exec(t_adr, scr)
    return buf

